# Comparator class to manage comparison 
# Takes as paramters, 2 lists of md5s
# Attributes for comparison:
# 1. Host name
# 2. IP location
# 3. Network file
# 4. Threat name
# 5. Other trait objects (e.g. API, String, Registry)

# library imports
from threatnamecompDB import ThreatNameCompDB
from networkcompDB import NetworkCompDB
from threatname import ThreatName
from trait import Trait
import threading
import ngram  # @UnresolvedImport
import re
import json

class Comparator:
    # initialise parameters
    def __init__(self, md5ListA, md5ListB):
        # handles the validated md5s in both list
        self.listA = self._validateAll(md5ListA)
        self.listB = self._validateAll(md5ListB)
        
        # handles the list of Threat(md5)
        self.threatListA = []
        self.threatListB = []
        
        # handles the list of Trait(md5)
        self.traitListA = []
        self.traitListB = []
        
        # handles the host, IP location comparison results
        # stores as dict
        self.networkHostResult = {}
        self.networkIPLocationResult = {}
        self.networkFileResult = {}
        
        # handles the details of Theat name comparison
        self.noOfThreatNameSimilarity = 0
        self.noOfThreatInListA = 0
        self.noOfThreatInListB = 0
        
        # handles the details of Network Host, IP location and Network files comparison
        self.noOfNetworkHostSimilarity = 0
        self.NoOfNetworkHostInListA = 0
        self.NoOfNetworkHostInListB = 0
        
        self.noOfNetworkIPLocationSimilarity = 0
        self.noOfNetworkIPLocationInListA = 0
        self.noOfNetworkIPLocationInListB = 0
        
        self.noOfNetworkFileSimilarity = 0
        self.noOfNetworkFileInListA = 0
        self.noOfNetworkFileInListB = 0
        
        # handles the details of trait comparison
        self.noOfTraitSimilarity = 0
        self.noOfTraitInListA = 0
        self.noOfTraitInListB = 0
        
        # handles the details of API comparison
        self.traitObjectNoDictInListA = 0
        self.traitObjectNoDictInListB = 0
        
        # executes thread for faster processing
        self._multiThread()

    # method to run multiple threads for Trait, Threat and Network comparison
    # Network initialisation and comparison uses a single thread
    # Trait initialisation uses a thread per MD5
    # Threat initialisation uses a thread per MD5
    def _multiThread(self):
        t4 = threading.Thread(target=self._getNetworkCompRes())
        t4.start()
        i = 0
        for md5 in self.listA:
            i = i+1
            t = threading.Thread(target=self._getThreatA, args=(md5,))
            t.start()
            t1 = threading.Thread(target=self._getTraitA, args=(md5,))
            t1.start()
        t.join()
        t1.join()
        for md5 in self.listB:
            t2 = threading.Thread(target=self._getThreatB, args=(md5,))
            t2.start()
            t3 = threading.Thread(target=self._getTraitB, args=(md5,))
            t3.start()
        t2.join()
        t3.join()
        t4.join()
    
    # helper method to create ThreatName object for list A
    def _getThreatA(self, md5):
        self.threatListA.append(ThreatName(md5))
        
    # helper method to create ThreatName object for list B
    def _getThreatB(self, md5):
        self.threatListB.append(ThreatName(md5))
    
    # helper method to create Trait object for list A
    def _getTraitA(self, md5):
        self.traitListA.append(Trait(md5))

    # helper method to create Trait object for list B
    def _getTraitB(self, md5):
        self.traitListB.append(Trait(md5))
   
    # method to compare ThreatName between the two ThreatName object list
    def getThreatCompRes(self):
        result = ThreatNameCompDB()
        listOfThreatA = []
        listOfThreatB = []
        for itemA in self.threatListA:
            listOfThreatA.extend(itemA.threatList)
            for itemB in self.threatListB:
                listOfThreatB.extend(itemB.threatList)
                similarityOfThreatGuesses = itemA.compareThreatGuesses(itemB.threatList) 
                result.add(similarityOfThreatGuesses, itemA.md5, itemB.md5)

        self.noOfThreatNameSimilarity = result.noOfSimilarity
        self.noOfThreatInListA = len(set(listOfThreatA))
        self.noOfThreatInListB = len(set(listOfThreatB))
        return result.threatCompRes       

    # helper method to contain Host Similarity
    def _containNetworkHostSimilarity(self, similarityOfNetwork, networkA, networkB):
        res = {}
        for item in similarityOfNetwork:
            setDict = {}
            if networkA.hostMD5ActionDict.has_key(item):
                setDict['Set A'] = networkA.hostMD5ActionDict[item]
            if networkB.hostMD5ActionDict.has_key(item):
                setDict['Set B'] = networkB.hostMD5ActionDict[item]
            res[item] = setDict
#         print json.dumps(res, indent =2)
        return res

    # helper method to contain Network IP Location Similarity
    def _containNetworkIPLocationSimilarity(self, similarityOfIPLocation, networkA, networkB):
        res = {}
        for location in similarityOfIPLocation:
            setDict = {}
            if networkA.locationMD5IPDict.has_key(location):
                setDict['Set A'] = networkA.locationMD5IPDict[location]
            if networkB.locationMD5IPDict.has_key(location):
                setDict['Set B'] = networkB.locationMD5IPDict[location]
            res[location] = setDict
#         print json.dumps(res, indent =2)
        return res
    
    # helper method to contain Network File Similarity  
    def _containNetworkFileSimilarity(self, similarityOfNetworkFile, networkA, networkB):
        res = {}
        for filename in similarityOfNetworkFile:
            setDict = {}
            if networkA.networkFileMD5RawDict.has_key(filename):
                setDict['Set A'] = networkA.networkFileMD5RawDict[filename]
            if networkB.networkFileMD5RawDict.has_key(filename):
                setDict['Set B'] = networkB.networkFileMD5RawDict[filename]
            res[filename] = setDict
#         print json.dumps(res, indent =2)
        return res
    
    # private method to compare Network of two md5 lists
    # handles Host, IP location and network file comparison
    def _getNetworkCompRes(self):
        # initialise NetworkCompDB for both lists of MD5s
        networkA = NetworkCompDB(self.listA)
        networkB = NetworkCompDB(self.listB)
        
        # set the details for Host comparison details 
        similarityOfNetworkHost = networkA.compareHost(networkB.hostListOnly)  
        self.NoOfNetworkHostInListA = len(list(set(networkA.hostMD5ActionDict)))
        self.NoOfNetworkHostInListB = len(list(set(networkB.hostMD5ActionDict)))
        self.noOfNetworkHostSimilarity = len(similarityOfNetworkHost)
        self.networkHostResult = self._containNetworkHostSimilarity(similarityOfNetworkHost, networkA, networkB)
     
        # set the details for IP Location comparison details 
        similarityOfIPLocation = networkA.compareLocation(networkB.ipLocationListOnly)
        self.noOfNetworkIPLocationInListA = len(list(set(networkA.ipLocationListOnly)))
        self.noOfNetworkIPLocationInListB = len(list(set(networkB.ipLocationListOnly)))
        self.noOfNetworkIPLocationSimilarity = len(similarityOfIPLocation)
        self.networkIPLocationResult = self._containNetworkIPLocationSimilarity(similarityOfIPLocation, networkA, networkB)
#         print self.networkIPLocationResult
        
        # set the details for Network File comparison details
        similarityOfNetworkFile = networkA.compareNetworkFile(networkB.networkFileListOnly)
        self.noOfNetworkFileInListA = len(list(set(networkA.networkFileListOnly)))
        self.noOfNetworkFileInListB = len(list(set(networkB.networkFileListOnly)))
        self.noOfNetworkFileSimilarity = len(similarityOfNetworkFile)
        self.networkFileResult = self._containNetworkFileSimilarity(similarityOfNetworkFile, networkA, networkB)
#         print self.networkIPLocationResul    
    
    # Combines the MD5s with the same trait in to a list
    # [{'d079d52ef1fd7a57e472d57a09df7e89': {u'trait2_object': u'Zero',}}, 
    #  {'c0ba07c36a293bf1cd8ae4957553fb8f': {u'trait2_object': u'Zero', ... }}]
    def _getAllTraits(self, traitList):  
        allTraits = {}
        for traitObject in traitList:   
            traitMD5ActionDict = traitObject.traitmd5actionDict
            for trait in traitMD5ActionDict:               
                if allTraits.has_key(trait):
                    existingDict = allTraits[trait]
                    md5actionDict = dict(existingDict.items() + traitMD5ActionDict[trait].items())
                else:
                    md5actionDict = traitMD5ActionDict[trait]
                allTraits[trait] = md5actionDict
        return allTraits
      

    # returns as dict of key trait and value trait object
    # e.g. {example_traitname: example_traitobject, .... }
    def _getTraitObjectType(self, allTraits, traitsOnly):
        traitObjectTypeDict = {}
        for trait in traitsOnly:
            val = allTraits[trait].values()[0]
            if val.has_key('trait_object'):
                traitObjectTypeDict[trait] = val['trait_object']
            else:
                traitObjectTypeDict[trait] =  val['trait2_object']
        return traitObjectTypeDict
    
    def _getTraitObjectNoDictInListA(self, traitObjectTypeDictA ):
        traitObjectNoDict = {}
        for trait in traitObjectTypeDictA.values():
            count = 1
            if traitObjectNoDict.has_key(trait):
                count = traitObjectNoDict[trait] + 1
            traitObjectNoDict[trait] = count
        return traitObjectNoDict
    
    # all traits -> {0F1C1B20-AC59-3BB9-E151-39F369633FB8}': [{'3cc9971f0b36eb7224e631983ab9fe04': {'trait': u'Global\\{0F1C1B....
    # packedSim ->  [{'md5 from set B': [{'10d4dbdfd14338617e0539ee84b22bf3': {'trait2': u'api="DrvGetModule...}
    # returns traitObjectPackedSimDict
    # e.g {traitObject_name : {no 1 : [{packedsim,....}
    def getTraitCompRes(self):
        allTraitsA = self._getAllTraits(self.traitListA)
        allTraitsB = self._getAllTraits(self.traitListB)
        traitsOnlyA = allTraitsA.keys()
        traitsOnlyB = allTraitsB.keys()
        traitObjectTypeDictA = self._getTraitObjectType(allTraitsA, traitsOnlyA)
        traitObjectTypeDictB = self._getTraitObjectType(allTraitsB, traitsOnlyB)
        
        # get the traitObjectNoDictInListA for individual trait objects
        self.traitObjectNoDictInListA = self._getTraitObjectNoDictInListA(traitObjectTypeDictA)
        self.traitObjectNoDictInListB = self._getTraitObjectNoDictInListA(traitObjectTypeDictB)
                
        traitObjectPackedSimDict = {}
        self.noOfTraitInListA = len(traitsOnlyA)
        self.noOfTraitInListB = len(traitsOnlyB)
        G = ngram.NGram(set(traitsOnlyA), key = self._lower)
        totalSim = 0
        for targetTrait in traitsOnlyB:
            number = 1
            existingSimDict= {}
            noPackedSim = {}
            packedSim = {}
            targetTraitObject = traitObjectTypeDictB[targetTrait]
            similarity = G.search(self._lower(targetTrait), threshold = 0.1)
            # similarity -> [(item,0.1),(item2,0.5)]
            if similarity: 
                count = 0
                maxCount = len(similarity)
                for item in similarity:
                    simTraitName = item[0]
                    simTraitObject = traitObjectTypeDictA[simTraitName]
                    if simTraitObject == targetTraitObject:
                        break
                    count += 1
                if count < maxCount:
                    if traitObjectPackedSimDict.has_key(targetTraitObject):
                        existingSimDict = traitObjectPackedSimDict[targetTraitObject]
                        number = len(existingSimDict) + 1
                    packedSim['Set B'] = allTraitsB[targetTrait]
                    packedSim['Set A'] = allTraitsA[simTraitName]
                    packedSim['Match %'] = int(similarity[count][1] * 100)
                    noPackedSim[number] = packedSim
                    newDict = dict(existingSimDict.items() + noPackedSim.items())
                    totalSim += 1
                    
                    traitObjectPackedSimDict[targetTraitObject] = newDict 
        self.noOfTraitSimilarity =  totalSim
        return traitObjectPackedSimDict
    # helper method to validate md5s in the list  
    
    def getTraitComResStats(self, traitObjectPackedSimDict):
        traitObjectStatsDict = {}
        for traitObject in traitObjectPackedSimDict:
            traitObjectStats = {'Total Matches': 0, '> 50% matched':0, '< 50% matched':0}
            for it in traitObjectPackedSimDict[traitObject]:
                matchLevel = traitObjectPackedSimDict[traitObject][it]["Match %"]
                if matchLevel >= 50:
                    traitObjectStats['> 50% matched'] += 1
                else:
                    traitObjectStats['< 50% matched'] += 1
                traitObjectStats['Total Matches'] += 1
            traitObjectStatsDict[traitObject] = traitObjectStats
        return traitObjectStatsDict
                         
    def _validateAll(self, md5list):
        for item in md5list:
            if self._validate(item) is not True:
                md5list.remove(item)
                print (item + " is removed")
        return md5list
    # helper method to validate md5   
    def _validate(self,md5):
        if (re.match('^[\w-]+$', md5) is None or len(md5) != 32):
            return False
        return True
    # turn the alphabets to lower case
    def _lower(self, s):
        return re.sub(r'\W+', '', s).lower()